# -*- coding: utf-8 -*-
# !/usr/bin/env python
"""
-------------------------------------------------
   File Name：     datalist
   Description :   
   Author :       lth
   date：          2022/8/3
-------------------------------------------------
   Change Activity:
                   2022/8/3 18:22: create this script
-------------------------------------------------
"""
__author__ = 'lth'

import random

import numpy as np
from PIL import Image
from torch.utils.data import Dataset
from torchvision import transforms


class FaceParsingData(Dataset):
    def __init__(self, data, mode="train"):
        super(FaceParsingData, self).__init__()
        self.mode = mode
        self.data = data

        self.image_size = [512, 512]

        self.transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
        ])

        self.train_transform = transforms.Compose([
            transforms.GaussianBlur(5),
            transforms.ColorJitter(brightness=0.5, contrast=0.5, saturation=0.5, hue=0.5)
        ])
        self.randomCrop = RandomCrop(self.image_size)
        self.blockMask = BlockMask(self.image_size)
        self.crop = True
        self.mask = True

    def __getitem__(self, index):
        img = Image.open(self.data[index]).convert("RGB")
        img = img.resize(self.image_size, Image.BILINEAR)
        label = Image.open(self.data[index].replace("CelebA-HQ-img", "mask").replace("jpg", "png")).convert("P")
        if self.mode == "train":
            rotate_angle = random.randint(0, 360)
            img = img.rotate(rotate_angle, fillcolor=(128, 128, 128))
            label = label.rotate(rotate_angle, fillcolor=(0, 0, 0))
            if self.crop and random.random() > 0.5:
                img, label = self.randomCrop(img, label)
            if self.mask and random.random() > 0.5:
                img, label = self.blockMask(img, label)
            if random.random() > 0.5:
                img = self.train_transform(img)
        img = self.transform(img)
        label = np.array(label)

        return img, label

    def __len__(self):
        return len(self.data)


class RandomCrop:
    def __init__(self, image_size):
        self.image_size = image_size

    def __call__(self, image, label):
        start_point = random.randint(0, self.image_size[0] // 2)
        end_point = random.randint(self.image_size[0] // 2, self.image_size[0])

        image = image.crop((start_point, start_point, end_point, end_point))
        label = label.crop((start_point, start_point, end_point, end_point))

        image = image.resize(self.image_size)
        label = label.resize(self.image_size)
        return image, label


class BlockMask:
    def __init__(self, image_size):
        self.image_size = image_size

    def __call__(self, image, label):
        w = random.randint(self.image_size[0] // 5, self.image_size[0] // 4)
        h = random.randint(self.image_size[0] // 5, self.image_size[0] // 4)
        mask_image = Image.new("RGB", (w, h), (128, 128, 128))
        mask_label = Image.new("RGB", (w, h), (0, 0, 0))
        start_point_x = random.randint(0, self.image_size[0])
        start_point_y = random.randint(0, self.image_size[0])

        image.paste(mask_image, (start_point_x, start_point_y, start_point_x + w, start_point_y + h))
        label.paste(mask_label, (start_point_x, start_point_y, start_point_x + w, start_point_y + h))

        return image, label


if __name__ == "__main__":
    image = "E:/Datasets2/CelebAMask-HQ/CelebA-HQ-img/6.jpg"
    label = "E:/Datasets2/CelebAMask-HQ/mask/6.png"

    image = Image.open(image)
    image = image.resize((512, 512))
    label = Image.open(label)

    blockMask = RandomCrop([512, 512])
    image, label = blockMask(image, label)
    image.show()
    label.show()
